#! /usr/bin/env python3
# -*- coding:utf-8 -*-
import argparse
import os
import sys
import time
import random
import platform
import re

import subprocess


HASH_RETRY = 5

class HashCollisionException(Exception):
    def __init__(self, message="backup fails for hash collides, try changing your hostname."):
        super().__init__(message)


class MethodNotSupportException(Exception):
    def __init__(self, method, message="unsupport method '{}', available choices {}."):
        self.methods_supported = ["nfs", "local"]
        self.message = message.format(method, str(self.methods_supported))
        super().__init__(self.message)


class MethodConflictException(Exception):
    def __init__(self, method, url, message="method local and remote url should not be assigned at the same time."):
        self.message = message
        super().__init__(self.message)


class NotRootPathException(Exception):
    def __init__(self, path, message="path '{}' is not available, absolute path should be given for safety sake."):
        self.message = message.format(path)
        super().__init__(self.message)

def p64(n: int) -> str:
    table = "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ-_"
    s = bin(n)[2:][::-1]
    return "".join([table[int(s[i : i + 6][::-1], 2)] for i in range(0, len(s), 6)][::-1])

def check_output(cmd):
    output = subprocess.check_output(cmd, shell=True)
    if type(output) is str:
        return output.strip()
    elif sys.version_info[0] == 2 and (type(output) is unicode or type(output) is bytes):
        return output.encode('utf-8').strip()
    elif sys.version_info[0] == 3 and type(output) is bytes:
        return str(output, 'utf-8').strip()
    else:
        raise TypeError("Expected bytes or string, but got %s." % type(output))


def check_call(cmd):
    return subprocess.check_call(cmd, shell=True)


def build_nfs_url(url, path):
    return "".join(["nfs://", url, '/', path])

def recover_preset():
    grub2_migrear_prefix = r"title=[\",\']?MigReaR-"
    output = check_output("grubby --info=ALL")
    grub_entries = output.split("index=")

    for entry in grub_entries:
        if re.findall(grub2_migrear_prefix, entry):
            items = re.findall(r"title=(.*)\n", entry)
            if len(items) > 0:
                index = items[0][0]
                check_call("grub2-set-default %s" % (index))
                return

    raise TypeError("Can not found backup grubby entry with title=MigReaR-xxxx")


def backup(method, initramfs_path, initramfs_url, layout_path, layout_url, host_ip, excludes, hash=None):
    if hash is None:
        tmp_mountponit = "/tmp/mnt"
        if not os.path.exists(tmp_mountponit):
            os.makedirs(tmp_mountponit)
        check_call("mount -v -t nfs -o rw,noatime %s:%s %s" % (host_ip, layout_path, tmp_mountponit))
        for _ in range(HASH_RETRY):
            hash = p64(int(time.time() * 10 ** 7)) + p64(random.randint(0, 1<<12)).rjust(2, '0')[-12:]
            if not os.path.exists("%s/%s-%s" % (tmp_mountponit, platform.node(), hash)):
                os.makedirs("%s/%s-%s" % (tmp_mountponit, platform.node(), hash))
                check_call("umount " + tmp_mountponit)
                break
        else:
            check_call("umount " + tmp_mountponit)
            raise HashCollisionException()

    local_cfg = """
OUTPUT=RAMDISK
OUTPUT_URL=%s
BACKUP=NETFS
BACKUP_URL=%s
BACKUP_PROG_EXCLUDE=("\$\{BACKUP_PROG_EXCLUDE[@]\}" '/media' '/var/tmp' '/var/crash' %s '/tmp/*')
NETFS_KEEP_OLD_BACKUP_COPY=
NETFS_PREFIX="$HOSTNAME-%s"
RAMDISK_SUFFIX="%s"
""" % (initramfs_url, layout_url, " ".join(excludes), hash, hash)

    with open("/etc/rear/local.conf", "w") as f:
        f.write(local_cfg)

    backup_env = " "

    backup_env = "MBACKUP_MODE=%s MBACKUP_INITRAMFS_PATH=%s MBACKUP_LAYOUT_PATH=%s " % (method.upper(), initramfs_path, layout_path)
    if method == "local":
        if not os.path.exists(layout_path):
            os.mkdir(layout_path)
        with open("/etc/exports", "a") as f:
            f.write("\n%s *(fsid=0,rw,sync,no_root_squash,no_subtree_check,crossmnt)" % (layout_path))
        check_call("systemctl restart nfs-server")

    command = " ".join([backup_env, "rear -v mkbackup"])
    check_call(command)

    # grub
    vmlinuz = check_output("ls /boot/vmlinuz-$(uname -r)").strip()

    initramfs_file_path = check_output("find %s -name 'initramfs-%s.img'" % (initramfs_path, hash))
    command = "grubby --copy-default --add-kernel %s --initrd %s --title MigReaR-%s " % (vmlinuz, initramfs_file_path, hash)
    if method == "nfs":
        args = " --args='ro console=tty0 console=ttyS0,115200n8 console=ttyS0 unattended'"
    elif method == "local":
        args = " --args='ro console=tty0 console=ttyS0,115200n8 console=ttyS0'"

    command += args
    check_call(command)


def parse_args():
    parser = argparse.ArgumentParser(description="Relex and Recover for OS migration.")
    parser.add_argument('--method', default=None, help="backup up method, local and nfs are supported currently.")
    parser.add_argument('--url', default=None, help="remote backup url")
    parser.add_argument('--excludes', default=None, help="excludes these directories or files when backup, splitted by comma.")
    parser.add_argument('--path', default='/storage', help="where root layout would be kept, /storage by default.")
    parser.add_argument('--boot', default=None, help="where initramfs would be kept, same with --path if not given.")
    parser.add_argument('--recover_preset', default=None, help="set recover")
    parser.add_argument('--hash', default=None, help="specific backup hash, if not given, the program will generate one to determine backup files among several machines.")

    args = parser.parse_args()

    args.excludes = args.excludes.split(',') if args.excludes else []

    if args.method == "local" and args.url is not None:
        raise MethodConflictException(args.method, args.url)
    if not args.boot:
        if args.method == "local":
            args.boot = args.path
        else:
            args.boot = "/boot/"

    for path in [args.path, args.boot] + args.excludes:
        if not path.startswith('/'):
            raise NotRootPathException(path)

    if args.method == "local":
        args.url = "127.0.0.1"
        args.excludes += [args.path]

    if args.method == "local":
        args.layout_path = build_nfs_url(args.url, args.path)
    elif args.method == "nfs":
        args.layout_path = build_nfs_url(args.url, args.path)

    args.initramfs_path = "file://" + args.boot

    return args

if __name__ == "__main__":
    args = parse_args()
    if args.recover_preset:
        recover_preset()
    else:
        backup(args.method, args.boot, args.initramfs_path, args.path, args.layout_path, args.url, args.excludes, args.hash)
